package com.jzsec.proxy.jzsecproxy.limit;

import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.env.Environment;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;
import redis.clients.jedis.exceptions.JedisConnectionException;

import javax.annotation.PostConstruct;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Date;

@Component
@Slf4j
public class RedisLimit {

	private String requestLimitNum = "1800";    // 每分钟产生的令牌总数
	@Autowired
	private RedisTemplate redisTemplate;
	@Autowired
	private Environment env;
	@PostConstruct
	private void init() {
		requestLimitNum = env.getProperty("request.limit.num", requestLimitNum);
	}

	public boolean acquire() {
		String luaScript =
				"local key = KEYS[1] \n" +
				"local limit = tonumber(ARGV[1]) \n" +
				"local current = tonumber(redis.call('get', key) or 0) \n" +
				"if current + 1 > limit then \n" +
				"    return 0 \n" +
				"else \n" +
				"    redis.call('INCRBY', key, '1'); \n" +    //  每次取 1个令牌
				"    redis.call('expire', key, '120') \n" +    // 120s 以后 redis key 过期
				"    return 1 \n" +
				"end \n";
		DefaultRedisScript<Integer> redisScript = new DefaultRedisScript<Integer>();
		redisScript.setScriptText(luaScript);
		redisScript.setResultType(Integer.class);
		String key = getMinuteKey();
		Object token = null;
		try {
			token = redisTemplate.execute(redisScript, new ArrayList<String>(){{add(key);}}, requestLimitNum);
		} catch (Exception e) {
			log.error("redis execute error", e);
			token = new String("0");
		}
		log.info("requestLimitNum = " + requestLimitNum + ", key = " + key + ", token = " + token.toString());
		return "1".equals(token.toString());  // 1-表示拿到令牌， 0-没有拿到令牌
	}

	public String getMinuteKey() {
		SimpleDateFormat sdf = new SimpleDateFormat("yyyyMMddHHmm");
		String res = sdf.format(new Date());
		return "LeanCloudLimit" + res;
	}

	public static void main(String[] args) {
		System.out.println(new RedisLimit().getMinuteKey());
	}
}
